/*
 * Copyright (C) 2018 Hein-Pieter van Braam
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <stdlib.h>
#include <stdio.h>
#include <stdarg.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>

#include "minizip/mz.h"
#include "minizip/mz_os.h"
#include "minizip/mz_zip.h"
#include "minizip/mz_strm.h"
#include "minizip/mz_strm_mem.h"

#include "ResStringPool.h"
#include "android-xml.h"

struct options_t {
	char* jarsigner = NULL;
	char* keystore = NULL;
	char* storepass = NULL;
	char* keyalias = NULL;
} options;

#ifndef __unix__
#define popen _popen
#endif

void print_error(const char* format, ...) {
	va_list args;
	va_start(args, format);
	int message_size = vsnprintf(NULL, 0, format, args);
	va_end(args);

	char* message = (char*) malloc(message_size + 1);
	va_start(args, format);
	vsnprintf(message, message_size + 1, format, args);
	va_end(args);

	perror(message);
	free(message);
	exit(1);
}

char* find_jarsigner() {
#ifdef __unix
	char output[512];
	const char* command = "/usr/bin/env which jarsigner";
	FILE* pipe = popen(command, "r");

	if (!pipe)
		print_error("/usr/bin/which not found");

	size_t size = fread(output, 1, 511, pipe);
	int retval = pclose(pipe);
	output[sizeof(output) - 1] = 0;

	if (retval == 0) {
		// strip newline
		output[strlen(output) - 1] = 0;
		return strdup(output);
	}
#else
	FILE* temp = fopen("C:\\Program Files\\Android\\Android Studio\\jre\\bin\\jarsigner.exe", "r");
	if (temp) {
		fclose(temp);
		return strdup("C:\\Program Files\\Android\\Android Studio\\jre\\bin\\jarsigner.exe");
	}
#endif
	printf("Could not locate jarsigner in $PATH\n");
	return NULL;
}

#ifdef __unix__
#define CMD "\"%s\" -verbose -sigalg SHA1withRSA -digestalg SHA1 -keystore \"%s\" -storepass \"%s\" \"%s\" \"%s\""
#else
#define CMD "\"\"%s\" -verbose -sigalg SHA1withRSA -digestalg SHA1 -keystore \"%s\" -storepass \"%s\" \"%s\" \"%s\"\""
#endif

void resign_apk(const char* apk) {
	char output[512];
	char* command;
	uint32_t command_size = snprintf(NULL, 0,
	CMD, options.jarsigner, options.keystore, options.storepass, apk,
			options.keyalias);

	command = (char*) malloc(command_size + 1);
	if (!command)
		print_error("Out of memory");

	snprintf(command, command_size + 1,
	CMD, options.jarsigner, options.keystore, options.storepass, apk,
			options.keyalias);

	printf("Running %s\n", command);
	int retval = system(command);
	if (retval) {
		printf("jarsigner failed\n");
		exit(retval);
	}

	free(command);
}

void print_help(const char* progname) {
	printf(
			"Usage: %s -j </path/to/jarsigner> -k <keystore> -p <keystore password> -a <keystore alias> apkfile.apk\n",
			progname);
	printf("\n");
	printf(
			"  example: %s -j /usr/bin/jarsigner -k ~/debug.keystore -p android -a androiddebugkey game.apk\n",
			progname);
	printf("\n");
#if __unix__
	printf(
			"  note that you don't have to supply jarsigner path if it is in $PATH\n");
#endif
	exit(1);
}

void fix_manifest(uint8_t* manifest_data, uint32_t manifest_size,
		uint8_t** fixed_manifest, uint32_t* fixed_manifest_size) {

	uint32_t file_location = 0;
	uint32_t out_size = 0;
	uint8_t *out;

	ResChunk_header main_header;
	memcpy(&main_header, manifest_data, sizeof(ResChunk_header));
	if (main_header.type != RES_XML_TYPE) {
		fprintf(stderr, "Not an Android binary XML file\n");
		exit(1);
	}

	if (main_header.size != manifest_size) {
		fprintf(stderr, "Header size does not match filesize\n");
		exit(1);
	}

	out_size = sizeof(ResChunk_header);
	out = (uint8_t*) malloc(out_size);
	if (!out) {
		print_error("Out of memory");
	}

	file_location += sizeof(ResChunk_header);
	ResChunk_header cur_header;

	ResStringPool* stringpool = NULL;

	bool skip_end = false;
	while (file_location < manifest_size) {
		memcpy(&cur_header, manifest_data + file_location,
				sizeof(ResChunk_header));

		if (cur_header.type == RES_XML_START_ELEMENT_TYPE) {
			if (!stringpool) {
				fprintf(stderr,
						"StringPool not yet initialized, bailing out\n");
				exit(1);
			}

			bool skip = false;

			ResXMLTree_node xml_header;
			memcpy(&xml_header, manifest_data + file_location,
					sizeof(ResXMLTree_node));
			ResXMLTree_attrExt attrs;
			memcpy(&attrs,
					manifest_data + file_location
							+ xml_header.header.headerSize,
					sizeof(ResXMLTree_attrExt));

			const char* name = stringpool->get(attrs.name.index);

			for (int i = 0; i < attrs.attributeCount; ++i) {
				ResXMLTree_attribute attribute;
				memcpy(&attribute,
						manifest_data + file_location
								+ xml_header.header.headerSize
								+ sizeof(ResXMLTree_attrExt)
								+ (i * attrs.attributeSize),
						sizeof(ResXMLTree_attribute));

				if (strcmp(name, "uses-permission") == 0) {
					const char* value = stringpool->get(
							attribute.rawValue.index);
					if (!(strncmp(value, "android.", 7) == 0)
							|| (strncmp(value, "com.android", 11) == 0)) {
						printf("Removing invalid permission '%s'\n", value);
						skip = true;
					}
				}
			}

			skip_end = skip;
			if (skip) {
				file_location += cur_header.size;
				continue;
			}
		}

		if (cur_header.type == RES_XML_END_ELEMENT_TYPE) {
			if (skip_end) {
				skip_end = false;
				file_location += cur_header.size;
				continue;
			}
		}

		if (cur_header.type == RES_STRING_POOL_TYPE) {
			stringpool = new ResStringPool(manifest_data + file_location);
		}

		out = (uint8_t*) realloc(out, out_size + cur_header.size);
		if (!out) {
			print_error("Out of memory");
		}
		memcpy(out + out_size, manifest_data + file_location, cur_header.size);
		out_size += cur_header.size;

		file_location += cur_header.size;
	}

	main_header.size = out_size;
	memcpy(out, &main_header, sizeof(ResChunk_header));

	*fixed_manifest = out;
	*fixed_manifest_size = out_size;
	delete stringpool;
}

int main(int argc, char* argv[]) {
	int c;

	while ((c = getopt(argc, argv, "j:k:p:a:")) != -1) {
		switch (c) {
		case 'j':
			options.jarsigner = strdup(optarg);
			break;
		case 'k':
			options.keystore = strdup(optarg);
			break;
		case 'p':
			options.storepass = strdup(optarg);
			break;
		case 'a':
			options.keyalias = strdup(optarg);
			break;
		case '?':
			print_help(argv[0]);
			break;
		default:
			print_help(argv[0]);
		}
	}

	const char* infile = NULL;
	for (int index = optind; index < argc; index++)
		infile = argv[index];

	if (!options.jarsigner) {
		options.jarsigner = find_jarsigner();
	}

	if (!options.jarsigner || !options.keyalias || !options.keystore
			|| !options.storepass || !infile) {
		print_help(argv[0]);
	}

	FILE* in = fopen(infile, "rb");
	if (!in)
		print_error("Error opening %s", infile);

	if (fseek(in, 0, SEEK_END) == -1)
		print_error("Fseek failed on %s", infile);

	long filesize = ftell(in);
	rewind(in);

	uint8_t* apk_data = (uint8_t*) malloc(filesize);
	if (!apk_data)
		print_error("Out of memory");

	fread(apk_data, filesize, 1, in);
	if (ferror(in))
		print_error("Error reading from %s", infile);

	fclose(in);

	char* outfile;
	uint32_t infile_len = strlen(infile);
	uint32_t dot = 0;

	for (int i = infile_len; i >= 0; --i) {
		if (infile[i] == '.') {
			dot = i;
			break;
		}
	}

	if (dot == 0)
		dot = infile_len;

	outfile = (char*) calloc(1, infile_len + 8);
	if (!outfile)
		print_error("Out of memory");

	memcpy(outfile, infile, dot);
	memcpy(outfile + dot, "_fixed", 6);
	memcpy(outfile + dot + 6, infile + dot, infile_len - dot);

	void* mem_stream = NULL;
	uint32_t err;

	mz_stream_mem_create(&mem_stream);
	mz_stream_mem_set_buffer(mem_stream, apk_data, filesize);
	mz_stream_open(mem_stream, NULL, MZ_OPEN_MODE_READ);

	void* in_zip_handle = mz_zip_open(mem_stream, MZ_OPEN_MODE_READ);
	err = mz_zip_locate_entry(in_zip_handle, "AndroidManifest.xml", NULL);
	if (err != MZ_OK) {
		fprintf(stderr, "File AndroidManifest.xml not found in the zip file\n");
		return err;
	}

	mz_zip_file* manifest_file_info = NULL;
	err = mz_zip_entry_get_info(in_zip_handle, &manifest_file_info);
	if (err != MZ_OK) {
		fprintf(stderr, "Error %d getting entry info in zip file\n", err);
		return err;
	}

	err = mz_zip_entry_read_open(in_zip_handle, 0, NULL);
	if (err != MZ_OK) {
		printf("Error %d opening entry in zip file, password protected file?\n",
				err);
		return err;
	}

	uint32_t manifest_size = manifest_file_info->uncompressed_size;
	uint8_t* manifest_data = (uint8_t*) malloc(manifest_size);
	if (!manifest_data) {
		print_error("Out of memory");
	}

	void *manifest_stream = NULL;

	mz_stream_mem_create(&manifest_stream);
	mz_stream_mem_set_buffer(manifest_stream, manifest_data, manifest_size);
	mz_stream_open(manifest_stream, NULL, MZ_OPEN_MODE_WRITE);

	uint8_t buf[INT16_MAX];
	int32_t read = 0;
	int32_t written = 0;
	while (1) {
		read = mz_zip_entry_read(in_zip_handle, buf, sizeof(buf));
		if (read < 0) {
			err = read;
			printf("Error %d reading entry in zip file\n", err);
			break;
		}

		if (read == 0)
			break;

		written = mz_stream_mem_write(manifest_stream, buf, read);

		if (written != read) {
			err = mz_stream_mem_error(manifest_stream);
			printf("Error %d in writing extracted file\n", err);
			break;
		}
	}

	mz_zip_entry_close(in_zip_handle);

	if (err != MZ_OK) {
		fprintf(stderr, "Failed to extract AndroidManifest.xml\n");
		return err;
	}

	uint8_t* fixed_manifest;
	uint32_t fixed_manifest_size;
	fix_manifest(manifest_data, manifest_size, &fixed_manifest,
			&fixed_manifest_size);

	printf("Fixed AndroidManifest.xml, old size: %i, new size %i\n",
			manifest_size, fixed_manifest_size);

	printf("Writing new APK %s\n", outfile);

	void* out_stream;
	mz_stream_os_create(&out_stream);
	err = mz_stream_os_open(out_stream, outfile, MZ_OPEN_MODE_CREATE);
	if (err != MZ_OK) {
		printf("Error opening file %s\n", outfile);
		exit(1);
	}

	void* out_zip_handle = mz_zip_open(out_stream, MZ_OPEN_MODE_WRITE);
	mz_zip_set_version_madeby(out_zip_handle, MZ_VERSION_MADEBY);

	err = mz_zip_goto_first_entry(in_zip_handle);
	while (err == MZ_OK) {
		mz_zip_file *file_info = NULL;
		uint8_t buf[INT16_MAX];
		int32_t read = 0;
		int32_t written = 0;

		err = mz_zip_entry_get_info(in_zip_handle, &file_info);
		if (err != MZ_OK) {
			printf("Error %d getting entry info in zip file\n", err);
			return err;
		}

		if (strncmp("META-INF", file_info->filename, 8) == 0) {
			err = mz_zip_goto_next_entry(in_zip_handle);
			continue;
		}

		if (strcmp("AndroidManifest.xml", file_info->filename) == 0) {
			printf("Adding fixed AndroidManifest.xml\n");

			file_info->uncompressed_size = fixed_manifest_size;

			err = mz_zip_entry_write_open(out_zip_handle, file_info,
			MZ_COMPRESS_LEVEL_DEFAULT, NULL);

			if (err == MZ_OK) {
				written = mz_zip_entry_write(out_zip_handle, fixed_manifest,
						fixed_manifest_size);

				if (written != fixed_manifest_size) {
					printf("Error opening %s for writing in new apk\n",
							file_info->filename);
				}
			} else {
				printf("Error opening %s for writing in new apk\n",
						file_info->filename);
			}

			mz_zip_entry_close(out_zip_handle);
			err = mz_zip_goto_next_entry(in_zip_handle);
			continue;
		}

		err = mz_zip_entry_read_open(in_zip_handle, 0, NULL);
		if (err != MZ_OK) {
			printf("Error %d opening entry in zip file\n", err);
			return err;
		}

		printf("Adding %s (%i bytes)\n", file_info->filename,
				file_info->uncompressed_size);

		err = mz_zip_entry_write_open(out_zip_handle, file_info,
		MZ_COMPRESS_LEVEL_DEFAULT, NULL);
		if (err == MZ_OK) {
			while (1) {
				read = mz_zip_entry_read(in_zip_handle, buf, sizeof(buf));
				if (read < 0) {
					err = read;
					printf("Error %d reading entry in zip file\n", err);
					break;
				}

				if (read == 0)
					break;

				written = mz_zip_entry_write(out_zip_handle, buf, read);

				if (written != read) {
					printf("Error in writing extracted file %i != %i\n",
							written, read);
					break;
				}
			}
		} else {
			printf("Error opening %s for writing in new apk\n",
					file_info->filename);
		}

		mz_zip_entry_close(in_zip_handle);
		mz_zip_entry_close(out_zip_handle);

		err = mz_zip_goto_next_entry(in_zip_handle);
	}

	mz_zip_close(in_zip_handle);
	mz_zip_close(out_zip_handle);
	mz_stream_mem_delete(&manifest_stream);
	mz_stream_mem_delete(&mem_stream);
	mz_stream_os_close(out_stream);
	mz_stream_os_delete(&out_stream);

	printf("Resigning %s\n", outfile);
	resign_apk(outfile);

	free(apk_data);
	free(manifest_data);
	free(fixed_manifest);
	free(outfile);

	free(options.jarsigner);
	free(options.keyalias);
	free(options.keystore);
	free(options.storepass);
}
